# 协调Agent - 负责协调其他Agent之间的协作
from python_a2a import A2AServer, run_server, TaskStatus, TaskState, AgentCard, AgentSkill, Task
import requests
import json
import re
import logging
from openai import OpenAI
import os
from dotenv import load_dotenv
from uuid import uuid4

# 设置日志
logger = logging.getLogger(__name__)

load_dotenv()

class CoordinatorAgent(A2AServer):
    def __init__(self, agent_card):
        super().__init__(agent_card=agent_card)
        # 其他Agent的地址
        self.agents = {
            "weather": "http://localhost:7003",
            "travel": "http://localhost:7004", 
            "restaurant": "http://localhost:7005",
            "budget": "http://localhost:7007"
        }
        
        # LLM客户端用于理解用户意图
        self.openai_client = OpenAI(
            api_key=os.getenv("DEEPSEEK_API_KEY"),
            base_url=os.getenv("DEEPSEEK_BASE_URL", "https://api.deepseek.com/v1"),
        )
        
        print(f"🎯 CoordinatorAgent 初始化完成")
        print(f"📍 已注册的Agent: {list(self.agents.keys())}")

    def handle_task(self, task):
        """处理用户任务，协调多个Agent完成复杂任务"""
        print("🎯 handle_task 方法被调用")
        logger.info("handle_task 方法被调用")
        
        user_query = task.message.content.text
        print(f"🎯 CoordinatorAgent收到用户请求: {user_query}")
        logger.info(f"CoordinatorAgent收到用户请求: {user_query}")
        
        # 生成会话ID
        session_id = str(uuid.uuid4())
        print(f"� 会话ID: {session_id}")
        
        # 1. 分析用户意图
        intent_analysis = self._analyze_user_intent(user_query)
        print(f"🧠 意图分析: {intent_analysis}")
        logger.info(f"意图分析结果: {intent_analysis}")
        
        # 2. 制定执行计划
        execution_plan = self._create_execution_plan(intent_analysis)
        print(f"📋 执行计划: {execution_plan}")
        logger.info(f"执行计划: {execution_plan}")
        
        # 推送计划创建事件
        print("📤 准备推送 plan_created 事件")
        self._emit_event("plan_created", {
            "session_id": session_id,
            "user_query": user_query,
            "intent": intent_analysis,
            "steps": execution_plan.get("steps", [])
        })
        print("✅ plan_created 事件已推送")
        
        # 3. 按顺序执行Agent调用
        results = {}
        final_response = ""
        
        try:
            for idx, step in enumerate(execution_plan["steps"], start=1):
                agent_name = step["agent"]
                agent_input = step["input"]
                
                print(f"🔄 执行步骤: 调用{agent_name}Agent")
                print(f"   输入数据: {json.dumps(agent_input, ensure_ascii=False, indent=2)}")
                logger.info(f"执行步骤: 调用{agent_name}Agent，输入: {agent_input}")
                
                # 发送步骤开始事件
                self._emit_event("step_start", {
                    "session_id": session_id,
                    "step_index": idx,
                    "agent": agent_name,
                    "input": agent_input,
                    "total_steps": len(execution_plan["steps"])
                })

                # 调用对应的Agent
                if agent_name in self.agents:
                    print(f"   📞 调用Agent地址: {self.agents[agent_name]}")
                    result = self._call_agent(agent_name, agent_input)
                    print(f"   ✅ {agent_name}Agent返回结果: {json.dumps(result, ensure_ascii=False, indent=2) if isinstance(result, dict) else str(result)[:200]}...")
                    results[agent_name] = result
                    logger.info(f"{agent_name}Agent返回结果: {result}")

                    # 发送步骤结果事件
                    self._emit_event("step_result", {
                        "session_id": session_id,
                        "step_index": idx,
                        "agent": agent_name,
                        "result": result,
                        "progress": int(idx / max(1, len(execution_plan["steps"])) * 100)
                    })
                    
                    # 将上一步的结果传递给下一步
                    if len(execution_plan["steps"]) > 1:
                        self._update_next_step_input(execution_plan, step, result)
                else:
                    print(f"   ❌ 找不到Agent: {agent_name}")
                    logger.error(f"找不到Agent: {agent_name}")
            
            # 4. 整合所有结果
            self._emit_event("finalizing", {"session_id": session_id, "message": "正在整合多Agent结果..."})
            final_response = self._integrate_results(user_query, results, intent_analysis)
            self._emit_event("final_result", {
                "session_id": session_id,
                "final_response": final_response,
                "agents_used": list(results.keys())
            })
            
        except Exception as e:
            final_response = f"执行过程中出现错误: {str(e)}"
            print(f"❌ 执行错误: {e}")

        # 构建响应
        response_data = {
            "agent": "CoordinatorAgent",
            "user_query": user_query,
            "intent_analysis": intent_analysis,
            "execution_results": results,
            "final_response": final_response,
            "agents_used": list(results.keys())
        }

        task.artifacts = [{"parts": [{"type": "text", "text": final_response}]}]
        task.status = TaskStatus(state=TaskState.COMPLETED)
        
        print(f"🎯 CoordinatorAgent最终回复: {final_response[:100]}...")
        return task

    def _emit_event(self, event_type: str, data: dict):
        """向web后端推送协商过程事件，用于前端实时可视化"""
        try:
            payload = {
                "type": event_type,
                "session_id": data.get("session_id"),
                "step": data  # 与后端 /api/negotiation/event 兼容
            }
            requests.post(
                "http://localhost:8000/api/negotiation/event",
                json=payload,
                timeout=2
            )
        except Exception as e:
            # 静默失败，不影响主流程
            logger.debug(f"事件推送失败 {event_type}: {e}")

    def _analyze_user_intent(self, query):
        """使用LLM分析用户意图"""
        try:
            prompt = f"""
            分析用户查询的意图，提取关键信息：
            用户查询: "{query}"
            
            特别注意：
            1. 如果用户提到预算、钱、元、费用等词汇，needs_budget应为true
            2. 如果用户提到天气、户外、室内等词汇，needs_weather应为true  
            3. 如果用户提到旅行、景点、玩等词汇，needs_attractions应为true
            4. 提取具体的数字信息（预算金额、天数等）
            
            请以JSON格式返回分析结果：
            {{
                "primary_intent": "主要意图（travel_planning/weather_query/restaurant_recommendation等）",
                "city": "目标城市",
                "date": "时间信息（如果有）",
                "needs_weather": true/false,
                "needs_attractions": true/false,
                "needs_restaurants": true/false,
                "needs_budget": true/false,
                "budget_amount": "预算金额数字（如果有）",
                "duration": "旅行天数（如果有）",
                "special_requirements": "特殊要求"
            }}
            """
            
            response = self.openai_client.chat.completions.create(
                model="deepseek-chat",
                messages=[{"role": "user", "content": prompt}],
                max_tokens=500
            )
            
            content = response.choices[0].message.content
            # 提取JSON部分
            json_match = re.search(r'\{.*\}', content, re.DOTALL)
            if json_match:
                return json.loads(json_match.group())
            else:
                # 备用简单解析
                return self._simple_intent_analysis(query)
                
        except Exception as e:
            print(f"LLM意图分析失败: {e}")
            return self._simple_intent_analysis(query)

    def _simple_intent_analysis(self, query):
        """简单的意图分析备用方案"""
        query_lower = query.lower()
        
        # 提取城市
        cities = ["北京", "beijing", "上海", "shanghai", "广州", "深圳"]
        city = "beijing"  # 默认
        for c in cities:
            if c in query_lower:
                city = "beijing" if c in ["北京", "beijing"] else c
                break
        
        # 提取预算信息
        budget_keywords = ["预算", "钱", "元", "费用", "花费", "成本"]
        needs_budget = any(keyword in query for keyword in budget_keywords)
        
        # 提取预算数字
        import re
        budget_match = re.search(r'(\d+)\s*元', query)
        budget_amount = int(budget_match.group(1)) if budget_match else None
        
        # 提取天数
        days_match = re.search(r'(\d+)\s*天', query)
        duration = int(days_match.group(1)) if days_match else None
        
        return {
            "primary_intent": "travel_planning",
            "city": city,
            "date": "今天" if "今天" in query else "明天" if "明天" in query else None,
            "needs_weather": "天气" in query or "适合" in query or "户外" in query,
            "needs_attractions": "去哪" in query or "景点" in query or "玩" in query or "旅行" in query,
            "needs_restaurants": "吃" in query or "餐厅" in query or "美食" in query,
            "needs_budget": needs_budget or budget_amount is not None,
            "budget_amount": budget_amount,
            "duration": duration,
            "special_requirements": query
        }

    def _create_execution_plan(self, intent_analysis):
        """根据意图分析制定执行计划"""
        steps = []
        
        # 如果需要预算分析，优先进行预算检查
        needs_budget = intent_analysis.get("needs_budget", False)
        print(f"🔍 制定执行计划 - needs_budget: {needs_budget}")
        logger.info(f"制定执行计划 - needs_budget: {needs_budget}")
        
        if needs_budget:
            budget_input = {
                "type": "collaboration_request",
                "user_requirements": {
                    "destination": intent_analysis.get("city", "beijing"),
                    "budget": f"{intent_analysis.get('budget_amount')}元" if intent_analysis.get("budget_amount") else "未指定",
                    "duration": f"{intent_analysis.get('duration')}天" if intent_analysis.get("duration") else "未指定",
                    "preferences": self._extract_preferences(intent_analysis.get("special_requirements", ""))
                },
                "collaboration_context": {
                    "purpose": "budget_analysis_and_optimization",
                    "coordinator_request": True
                }
            }
            steps.append({
                "agent": "budget",
                "input": budget_input,
                "description": "分析预算可行性"
            })
            print(f"✅ 添加BudgetAgent到执行计划")
            logger.info(f"添加BudgetAgent到执行计划，输入: {budget_input}")
        
        # 如果需要天气信息，先查询天气
        if intent_analysis.get("needs_weather", False):
            steps.append({
                "agent": "weather",
                "input": {
                    "city": intent_analysis.get("city", "beijing"),
                    "query": f"查询{intent_analysis.get('city', 'beijing')}的天气"
                },
                "description": "获取天气信息"
            })
        
        # 如果需要景点推荐
        if intent_analysis.get("needs_attractions", False):
            steps.append({
                "agent": "travel",
                "input": {
                    "city": intent_analysis.get("city", "beijing"),
                    "weather_analysis": "来自weather_agent的结果"  # 占位符
                },
                "description": "获取景点推荐"
            })
        
        # 如果需要餐厅推荐
        if intent_analysis.get("needs_restaurants", False):
            steps.append({
                "agent": "restaurant", 
                "input": {
                    "city": intent_analysis.get("city", "beijing"),
                    "recommendations": "来自travel_agent的结果"  # 占位符
                },
                "description": "获取餐厅推荐"
            })
        
        return {"steps": steps, "intent": intent_analysis}

    def _call_agent(self, agent_name, input_data):
        """调用指定的Agent"""
        try:
            agent_url = self.agents[agent_name]
            endpoint = f"{agent_url}/a2a/tasks/send"
            
            payload = {
                "message": {
                    "content": {
                        "text": json.dumps(input_data, ensure_ascii=False)
                    }
                }
            }
            
            response = requests.post(endpoint, json=payload, timeout=30)
            response.raise_for_status()
            
            result = response.json()
            response_text = result["artifacts"][0]["parts"][0]["text"]
            
            # 尝试解析为JSON
            try:
                return json.loads(response_text)
            except:
                return {"raw_response": response_text}
                
        except Exception as e:
            print(f"调用{agent_name}Agent失败: {e}")
            return {"error": str(e)}

    def _update_next_step_input(self, execution_plan, current_step, result):
        """将当前步骤的结果传递给下一步"""
        current_index = execution_plan["steps"].index(current_step)
        
        if current_index < len(execution_plan["steps"]) - 1:
            next_step = execution_plan["steps"][current_index + 1]
            
            if current_step["agent"] == "weather" and next_step["agent"] == "travel":
                # 将天气结果传递给旅游Agent
                next_step["input"].update(result)
            elif current_step["agent"] == "travel" and next_step["agent"] == "restaurant":
                # 将景点结果传递给餐厅Agent
                next_step["input"].update(result)

    def _integrate_results(self, user_query, results, intent_analysis):
        """整合所有Agent的结果生成最终回复"""
        response_parts = []
        
        response_parts.append(f"针对您的问题「{user_query}」，我为您规划了完整的行程：\n")
        
        # 预算分析
        if "budget" in results:
            budget_data = results["budget"]
            print(f"🔍 预算数据结构: {json.dumps(budget_data, ensure_ascii=False, indent=2) if isinstance(budget_data, dict) else str(budget_data)}")
            
            try:
                # 直接处理JSON数据（BudgetAgent返回的格式）
                if isinstance(budget_data, dict) and "budget_data" in budget_data:
                    budget_info = budget_data["budget_data"]
                    status = budget_info.get("budget_status", "unknown")
                    
                    if status == "impossible":
                        response_parts.append(f"💰 **预算分析**：⚠️ 预算不足")
                        response_parts.append(f"   - 最低需求：{budget_info.get('min_required_budget', 'N/A')}元")
                        response_parts.append(f"   - 建议：{budget_info.get('recommendation', 'N/A')}")
                    elif status == "sufficient":
                        response_parts.append(f"💰 **预算分析**：✅ 预算充足")
                        breakdown = budget_info.get("optimized_breakdown", {})
                        total = budget_info.get("total_required", 0)
                        response_parts.append(f"   - 总预算：{total}元")
                        if breakdown:
                            response_parts.append(f"   - 预算分配：住宿 {breakdown.get('accommodation', 0)}元，餐饮 {breakdown.get('food', 0)}元")
                            response_parts.append(f"              交通 {breakdown.get('transportation', 0)}元，活动 {breakdown.get('activities', 0)}元")
                    else:
                        response_parts.append(f"💰 **预算分析**：状态 {status}")
                    
                    # 显示优化建议
                    if "impact_analysis" in budget_data:
                        impact = budget_data["impact_analysis"]
                        if "cost_optimization_potential" in impact:
                            opt = impact["cost_optimization_potential"]
                            savings = opt.get("potential_savings", 0)
                            if savings > 0:
                                response_parts.append(f"   - 💡 优化建议：可节省约 {savings}元")
                                
                # 如果是包装在artifacts中的格式（备用处理）  
                elif isinstance(budget_data, dict) and "artifacts" in budget_data:
                    artifact_text = budget_data["artifacts"][0]["parts"][0]["text"]
                    budget_analysis = json.loads(artifact_text)
                    
                    if "budget_data" in budget_analysis:
                        budget_info = budget_analysis["budget_data"]
                        status = budget_info.get("budget_status", "unknown")
                        
                        if status == "impossible":
                            response_parts.append(f"💰 **预算分析**：⚠️ 预算不足")
                        elif status == "sufficient":
                            response_parts.append(f"💰 **预算分析**：✅ 预算充足")
                        else:
                            response_parts.append(f"💰 **预算分析**：状态 {status}")
                            
            except Exception as e:
                response_parts.append(f"💰 **预算分析**：处理中... (错误: {str(e)})")
                print(f"❌ 预算数据解析错误: {e}")
        
        # 天气信息
        if "weather" in results:
            weather_data = results["weather"]
            if "weather_raw" in weather_data:
                response_parts.append(f"🌤️ **天气情况**：{weather_data['weather_raw']}")
                if "weather_analysis" in weather_data:
                    analysis = weather_data["weather_analysis"]
                    response_parts.append(f"   - 户外活动适宜性：{analysis['outdoor_activities']}")
                    if analysis.get("recommendations"):
                        response_parts.append(f"   - 建议活动类型：{', '.join(analysis['recommendations'])}")
        
        # 景点推荐
        if "travel" in results:
            travel_data = results["travel"]
            if "recommendations" in travel_data:
                response_parts.append(f"\n🏛️ **推荐景点**：")
                for i, attraction in enumerate(travel_data["recommendations"][:3], 1):
                    reason = attraction.get("reason", "")
                    note = attraction.get("note", "")
                    response_parts.append(f"   {i}. {attraction['name']} ({attraction['type']})")
                    if reason:
                        response_parts.append(f"      推荐理由：{reason}")
                    if note:
                        response_parts.append(f"      温馨提示：{note}")
        
        # 餐厅推荐
        if "restaurant" in results:
            restaurant_data = results["restaurant"]
            if "restaurant_recommendations" in restaurant_data:
                response_parts.append(f"\n🍽️ **美食推荐**：")
                for i, restaurant in enumerate(restaurant_data["restaurant_recommendations"][:3], 1):
                    location_reason = restaurant.get("location_reason", "")
                    response_parts.append(f"   {i}. {restaurant['name']} - {restaurant['type']} ({restaurant['price']})")
                    if location_reason:
                        response_parts.append(f"      位置优势：{location_reason}")
        
        response_parts.append(f"\n✨ 希望这个综合规划对您的{intent_analysis.get('city', '目的地')}之行有所帮助！")
        
        return "\n".join(response_parts)
    
    def _extract_preferences(self, query):
        """从查询中提取用户偏好"""
        preferences = []
        if "户外" in query:
            preferences.append("户外活动")
        if "美食" in query:
            preferences.append("美食体验")
        if "历史" in query or "文化" in query:
            preferences.append("历史文化")
        if "摄影" in query:
            preferences.append("摄影")
        if "购物" in query:
            preferences.append("购物")
        if "休闲" in query:
            preferences.append("休闲放松")
        return preferences if preferences else ["综合体验"]

if __name__ == "__main__":
    agent_card = AgentCard(
        name="Coordinator Agent",
        description="智能协调Agent，整合天气、景点、餐厅等多个专业Agent提供完整的旅行规划服务",
        url="http://localhost:7006",
        version="1.0.0",
        skills=[
            AgentSkill(
                name="Multi-Agent Coordination",
                description="协调多个专业Agent协同工作"
            ),
            AgentSkill(
                name="Intent Analysis",
                description="分析用户复杂需求并制定执行计划"
            ),
            AgentSkill(
                name="Result Integration", 
                description="整合多个Agent的结果生成综合回复"
            ),
            AgentSkill(
                name="Travel Planning",
                description="提供完整的旅行规划服务"
            )
        ]
    )

    coordinator_agent = CoordinatorAgent(agent_card)
    
    print("🎯 CoordinatorAgent 启动于 http://localhost:7006")
    print("🚀 多Agent协作系统已就绪！")
    run_server(coordinator_agent, host="0.0.0.0", port=7006)
